# -*- coding: utf-8 -*-
import cv2
import numpy as np

#使用绝对地址，避免出错
datapath = "train_images"  # 训练样本地址
SAMPLES = 40               # 训练样本数量
pos, neg = "pos/pos-", "nes/nes-"

'''
定义训练用图像地址
'''
def path(cls,i):
    return "%s/%s%d.jpg"  % (datapath,cls,i)  #注意图片的格式，不能写错

'''定义FLANN匹配器函数'''
def get_flann_matcher():
    flann_params = dict(algorithm = 1, trees = 5)
    return cv2.FlannBasedMatcher(flann_params, {})

'''定义SIFT特征检测器/提取器'''
def get_bow_extractor(extract, match):
    return cv2.BOWImgDescriptorExtractor(extract, match)

'''定义BOW训练器'''
def get_extract_detect():
    return cv2.xfeatures2d.SIFT_create(), cv2.xfeatures2d.SIFT_create()

'''用于检测并返回图像的特征'''
def extract_sift(fn, extractor, detector):
    im = cv2.imread(fn,0)
    return extractor.compute(im, detector.detect(im))[1]

'''提取BOW特征，并返回他们'''  
def bow_features(img, extractor_bow, detector):
    return extractor_bow.compute(img, detector.detect(img))

'''进行检测器训练'''
def car_detector():
    '''创建/导入训练用对象'''
#    pos, neg = "pos/pos-", "nes/nes-"
    detect, extract = get_extract_detect()
    matcher = get_flann_matcher()
    #extract_bow = get_bow_extractor(extract, matcher)
    print("building BOWKMeansTrainer...")
    bow_kmeans_trainer = cv2.BOWKMeansTrainer(12)
    extract_bow = cv2.BOWImgDescriptorExtractor(extract, matcher)

    """
    向训练器增加训练图像特征，
    对于每个类，将向训练器和负样本图像添加一个正样本图像，
    然后，训练器会将数据聚成k个类（k_means方法）
    """
    print("adding features to trainer")
    for i in range(SAMPLES):
        print(i)
        bow_kmeans_trainer.add(extract_sift(path(pos,i), extract, detect))
        #bow_kmeans_trainer.add(extract_sift(path(neg,i), extract, detect))
  
    '''设置BOWImgDescriptorExtractor类的词汇'''
    vocabulary = bow_kmeans_trainer.cluster()
    extract_bow.setVocabulary(vocabulary)

  
    """
    开始进行训练，
    将训练数据与类进行关联，
    含检测目标的用1表示，不含的用-1表示。
    """
    '''
    用两个数组分别表示训练数据集和训练标签，
    通过遍历数据集来填充这两个数组。
    这段代码每循环一次，
    都会增加一个正样本图像和一个负样本图像，
    并将他们的标签设为1和-1，以保证标签和样本一致。
    '''
    traindata, trainlabels = [],[]
    print("adding to train data")
    for i in range(SAMPLES):
        print(i)
        traindata.extend(bow_features(cv2.imread(path(pos, i), 0), extract_bow, detect))
        trainlabels.append(1)
        traindata.extend(bow_features(cv2.imread(path(neg, i), 0), extract_bow, detect))
        trainlabels.append(-1)
        """
        #若想设为更多的类，可以改为：
        traindata.extend(bow_features(cv2.imread(path(class1, i), 0), extract_bow, detect))
        trainlabels.append(1)
        traindata.extend(bow_features(cv2.imread(path(class2, i), 0), extract_bow, detect))
        trainlabels.append(2)
        indata.extend(bow_features(cv2.imread(path(class3, i), 0), extract_bow, detect))
        trainlabels.append(3)
        traindata.extend(bow_features(cv2.imread(path(class4, i), 0), extract_bow, detect))
        trainlabels.append(4)
        """

  
    '''
    训练一个SVM。
    '''
    svm = cv2.ml.SVM_create()
    svm.setType(cv2.ml.SVM_C_SVC)
    svm.setGamma(1)
    '''
    svm.setC的参数决定分类器的训练误差和预测误差，
    值越大，误判的可能性就越小，但训练的精度就越低，
    值越小就越容易过拟合，降低预测精度。 
    '''
    svm.setC(35) 
    '''
    svm.setKernel决定分类器的性质：
    SVM_LINEAR说明分类器为线性超平面，主要用于二分类；
    SVM_RBF使用高斯函数对数据进行分类，数据被分到这些函数的核中；
    当使用SVM分的类超过2个时就必须使用SVM_RBF。
    '''
    svm.setKernel(cv2.ml.SVM_RBF)

    '''将traindata, trainlabels传给SVM训练函数'''
    svm.train(np.array(traindata), cv2.ml.ROW_SAMPLE, np.array(trainlabels))
    #返回SVM、BOW提取器对象
    #svm.save('svmxml') #用来保存训练的模型
    return svm, extract_bow
